#!/usr/bin/env python
#
# Copyright 2005, 2006 Free Software Foundation, Inc.
# 
# This file is part of GNU Radio
# 
# GNU Radio is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 3, or (at your option)
# any later version.
# 
# GNU Radio is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
# 
# You should have received a copy of the GNU General Public License
# along with GNU Radio; see the file COPYING.  If not, write to
# the Free Software Foundation, Inc., 51 Franklin Street,
# Boston, MA 02110-1301, USA.

import copy, math, sys
from gnuradio import gr, gru, tx80211g, usrp2, eng_notation
import gnuradio.gr.gr_threading as _threading
from gnuradio.blks2impl import psk, qam
import tx80211g_packet_utils as ofdm_packet_utils

#transmit path
class tx80211g_transmit_path(gr.hier_block2):
    def __init__(self, options, payload=''):
	gr.hier_block2.__init__(self, "transmit_path",
				gr.io_signature(0, 0, 0), # Input signature
				gr.io_signature(1, 1, gr.sizeof_gr_complex)) # Output signature

	#modulator
	self.ofdm_tx       = ofdm_mod(options, payload, msgq_limit=2, pad_for_usrp=False)

	# static value to make sure we do not exceed +-1 for the floats being sent to the sink
	self._norm         = options.norm
	self.amp = gr.multiply_const_cc(self._norm)

    self.connect(self.ofdm_tx, self.amp, self)
        
    def send_pkt(self, payload='', eof=False):
        return self.ofdm_tx.send_pkt(payload, eof)

class ofdm_mod(gr.hier_block2):
    def __init__(self, options, payload='', msgq_limit=2, pad_for_usrp=False):
        """
        output-complex modulated signal

        options: modulation options from higher layers
        msgq_limit(int): max messages in message queue        
        pad_for_usrp: If true, packets are padded such that they end up a multiple of 128 samples
        """

	gr.hier_block2.__init__(self, "ofdm_mod",
				gr.io_signature(0, 0, 0),       # Input signature
				gr.io_signature(1, 1, gr.sizeof_gr_complex)) # Output signature
	
    
    self._fft_length          = 64
    self._total_sub_carriers  = 53
	self._data_subcarriers    = 48
	self._cp_length           = 16
 	self._regime              = options.regime
	self._symbol_length       = self._fft_length + self._cp_length
	
	self._symbol_time	  = options.interp*(self._fft_length+self._cp_length)/100
	
    win = []

	if(self._regime == "1" or self._regime == "2"):
	    rotated_const = ofdm_packet_utils.bpsk(self)
        
    elif (self._regime == "3" or self._regime == "4"):
	    rotated_const = ofdm_packet_utils.qpsk(self)

    elif(self._regime == "5" or self._regime == "6"):
        rotated_const = ofdm_packet_utils.qam16(self)

    elif(self._regime == "7" or self._regime == "8"):
        rotated_const = ofdm_packet_utils.qam64(self)
	
    #map groups of bits
    self._pkt_input = tx80211g.ofdm_mapper(rotated_const, msgq_limit, self._data_subcarriers, self._fft_length)
        
    #insert pilot symbols
    self.pilot = tx80211g.ofdm_pilot_cc(self._data_subcarriers)
	
    #move subcarriers to correct place
    self.cmap  = tx80211g.ofdm_cmap_cc(self._fft_length, self._total_sub_carriers)

    
	#inverse fast fourier transform
    self.ifft = gr.fft_vcc(self._fft_length, False, win, False)

    #cyclic prefix
    self.cp_adder = gr.ofdm_cyclic_prefixer(self._fft_length, self._symbol_length)
    
    #scale
    self.scale = gr.multiply_const_cc(1.0 / math.sqrt(self._fft_length))
        
	#pass number of OFDM data symbols for preamble and zerogap
	info = ofdm_packet_utils.get_info(payload, options.regime, self._symbol_time)
	N_sym             = info["N_sym"]
	
	#training sequence
    self.preamble= ofdm_packet_utils.insert_preamble(self._symbol_length, N_sym)

    #append zero samples for the receiver
	self.zerogap    = ofdm_packet_utils.insert_zerogap(self._symbol_length, N_sym)
        
    #repeat the frame (if option is set)
    self.repeat = tx80211g.repetition(80, options.repetition, N_sym)

	self.s2v = gr.stream_to_vector(gr.sizeof_gr_complex , self._symbol_length)
	self.v2s = gr.vector_to_stream(gr.sizeof_gr_complex , self._symbol_length)
	
	#swap real and immaginary component before sending (GNURadio/USRP2 bug!)
    #taken from ftw project
	if options.swapIQ == True:
		self.gr_complex_to_imag_0 = gr.complex_to_imag(1)
		self.gr_complex_to_real_0 = gr.complex_to_real(1)
		self.gr_float_to_complex_0 = gr.float_to_complex(1)
		self.connect((self.v2s, 0), (self.gr_complex_to_imag_0, 0))
		self.connect((self.v2s, 0), (self.gr_complex_to_real_0, 0))
		self.connect((self.gr_complex_to_real_0, 0), (self.gr_float_to_complex_0, 1))
		self.connect((self.gr_complex_to_imag_0, 0), (self.gr_float_to_complex_0, 0))
		self.connect((self.gr_float_to_complex_0, 0), (self))
	elif options.swapIQ == False:
		self.gr_complex_to_imag_0 = gr.complex_to_imag(1)
		self.gr_complex_to_real_0 = gr.complex_to_real(1)
		self.gr_float_to_complex_0 = gr.float_to_complex(1)
		self.connect((self.v2s, 0), (self.gr_complex_to_imag_0, 0))
		self.connect((self.v2s, 0), (self.gr_complex_to_real_0, 0))
		self.connect((self.gr_complex_to_imag_0, 0), (self.gr_float_to_complex_0, 1))
		self.connect((self.gr_complex_to_real_0, 0), (self.gr_float_to_complex_0, 0))
		self.connect((self.gr_float_to_complex_0, 0), (self))
		
    #connect all the blocks
	self.connect((self._pkt_input, 0), (self.pilot, 0))
	self.connect((self._pkt_input,1), (self.preamble, 1))
	self.connect((self.preamble,1), (self.zerogap, 1))
	
	self.connect(self.pilot, self.cmap, self.ifft, self.cp_adder, self.scale, self.s2v, self.preamble, self.zerogap, self.repeat, self.v2s)
    def send_pkt(self, payload='', eof=False):
        """        
        payload (string): data to send      
        """
        if eof:
            msg = gr.message(1)
        else:
            info = ofdm_packet_utils.get_info(payload, self._regime, self._symbol_time)
	    N_cbps            = info["N_cbps"]
	    N_bpsc            = info["N_bpsc"]
        N_rate            = info["rate"]
	    N_sym             = info["N_sym"]

	    (pkt,Length) = ofdm_packet_utils.tx80211g_make(payload,self._regime, self._symbol_time)
	    (pkt_scrambled,Length) = ofdm_packet_utils.scrambler(pkt,Length)
	    pkt_coded = ofdm_packet_utils.conv_encoder(pkt_scrambled, Length, self._regime, N_cbps, N_bpsc, N_sym, N_rate)
	    pkt_interleaved = ofdm_packet_utils.interleaver(pkt_coded , self._regime, N_cbps, N_bpsc)
	    msg = gr.message_from_string(pkt_interleaved)

        self._pkt_input.msgq().insert_tail(msg)

class tx80211g_receive_path(gr.hier_block2):
    def __init__(self, options):
        gr.hier_block2.__init__(self, "usrp_rx", gr.io_signature(0, 0, 0), gr.io_signature(1, 1, gr.sizeof_gr_complex))

        self.u = usrp2.source_32fc(options.interface, options.mac_addr)
        gain = options.gain
        freq = options.freq
        if gain is None:
            g = self.u.gain_range()
            gain = float(g[0]+g[1])/2

        if verbose:
            print "gain = ", gain

        if freq is None:
            r = self.u.freq_range()
            freq = float(r[0]+r[1])/2

        self.u.set_gain(gain)
        r = self.u.set_center_freq(freq)
        if verbose:
            print "desired freq = ", freq
        
        self.connect(self.u, self)    
        
#Experimental: not working as of now..        
class ofdm_demod(gr.hier_block2):    
    def __init__(self, options, callback=None):
        """
	    output-complex modulated signal

        options: pass modulation options
        msgq_limit(int): maximum number of messages in message queue
        pad_for_usrp: If true, packets are padded to make them multiple of 128
        """

	self._rcvd_pktq = gr.msg_queue()          #PHY packets


    self._watcher = _queue_watcher_thread(self._rcvd_pktq, callback)
	
    
    self._fft_length          = 64
    self._total_sub_carriers  = 53
	self._data_subcarriers    = 48
	self._cp_length           = 16
 	self._regime              = options.regime
	self._symbol_length       = self._fft_length + self._cp_length
		
	self._symbol_time	  = options.interp*(self._fft_length+self._cp_length)/100
	
    win = []

	if(self._regime == "1" or self._regime == "2"):
	    rotated_const = ofdm_packet_utils.bpsk(self)
        
    elif (self._regime == "3" or self._regime == "4"):
	    rotated_const = ofdm_packet_utils.qpsk(self)

    elif(self._regime == "5" or self._regime == "6"):
        rotated_const = ofdm_packet_utils.qam16(self)

    elif(self._regime == "7" or self._regime == "8"):
        rotated_const = ofdm_packet_utils.qam64(self)
	
    self._pkt_input = tx80211g.ofdm_mapper(rotated_const, msgq_limit, self._data_subcarriers, self._fft_length)
        
    self.pilot = tx80211g.ofdm_pilot_cc(self._data_subcarriers)
	    
    self.cmap  = tx80211g.ofdm_cmap_cc(self._fft_length, self._total_sub_carriers)

    
	self.ifft = gr.fft_vcc(self._fft_length, False, win, False)

    self.cp_adder = gr.ofdm_cyclic_prefixer(self._fft_length, self._symbol_length)
    
    self.scale = gr.multiply_const_cc(1.0 / math.sqrt(self._fft_length))
        
	info = ofdm_packet_utils.get_info(payload, options.regime, self._symbol_time)
	N_sym             = info["N_sym"]
	
	self.preamble= ofdm_packet_utils.insert_preamble(self._symbol_length, N_sym)

    self.zerogap    = ofdm_packet_utils.insert_zerogap(self._symbol_length, N_sym)
        
    self.repeat = tx80211g.repetition(80, options.repetition, N_sym)

	self.s2v = gr.stream_to_vector(gr.sizeof_gr_complex , self._symbol_length)
	self.v2s = gr.vector_to_stream(gr.sizeof_gr_complex , self._symbol_length)
	
    self.connect((self._pkt_input, 0), (self.pilot, 0))
	self.connect((self._pkt_input,1), (self.preamble, 1))
	self.connect((self.preamble,1), (self.zerogap, 1))
	
	self.connect(self.pilot, self.cmap, self.ifft, self.cp_adder, self.scale, self.s2v, self.preamble, self.zerogap, self.repeat, self.v2s)                